
import pandas as pd
import tushare as ts
from jqdatasdk import *
import time,datetime


# 单因子market_cap_top 市值最小  上证50轮动 模型
ts.set_token('883c1e2da92c0f68194fe2c4b4aa0e4da880b41ef62ff556f2da00e7')
pro = ts.pro_api()

auth('15068762322','keepstr123')
# 时间类型转换 如2013-01-01 -> 20130101


def charge_time(l=[]):
    l1 = list()
    for x in l:
        list_i = list(x)
        list_i.pop(4)
        list_i.pop(6)
        str_i = ''.join(list_i)

        l1.append(str_i)
    return l1


def charge_time1(date):

    list_i = list(date)
    list_i.insert(4,'-')
    list_i.insert(7,'-')
    str_i = ''.join(list_i)
    return str_i
def charge_time1(date):

    list_i = list(date)
    list_i.insert(4,'-')
    list_i.insert(7,'-')
    str_i = ''.join(list_i)
    return str_i
def getBeforeMonth(timestamp,months):#format='%Y-%m-%d %H:%M:%S'):
    '''
    以给定时间戳为基准，后退 months 个月后得到对应的时间戳
    '''
    from calendar import monthrange
    now_time=datetime.datetime.strptime(timestamp,'%Y-%m-%d')
    year,month,day=[int(one) for one in str(now_time).split(' ')[0].split('-')]
    for i in range(months):
        now_time-=datetime.timedelta(days=monthrange(year,month)[1])
    next_timestamp=now_time.strftime('%Y-%m-%d')
    tempvalue = next_timestamp.split("-")
    if tempvalue[1] in ['01', '02', '03']:
        return (tempvalue[0] + "-03-31")
    elif tempvalue[1] in ['04', '05', '06']:
        return (tempvalue[0] +  "-06-30")
    elif tempvalue[1] in ['07', '08', '09']:
        return (tempvalue[0] +  "-09-30")
    elif tempvalue[1] in ['10', '11', '12']:
        return (tempvalue[0] +  "-12-31")
    return next_timestamp
def get_data(ts_code='',start_date='',end_date='',retry_count=5,pause=5):
    for _ in range(retry_count):
        try:
            price = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
        except:
            time.sleep(pause)
        else:
            return price

def getcode(date):

    codeList = get_index_stocks('000016.XSHG', date=date)
    q = query(
        valuation.code, valuation.market_cap
    ).filter(
        valuation.code.in_(codeList)
    )
    df = get_fundamentals(q,getBeforeMonth(date,3))
    code = df[df.market_cap < 0]['code']
    df.set_index(['code'],inplace=True)
    for i in code.values.tolist():
        df.drop([i],axis=0,inplace=True)
    df.reset_index(inplace=True)
    df.sort_values(by='market_cap', axis=0, inplace=True)
    df = df.head(10)
    df['sum'] = round(df['market_cap'].sum(),4)
    df['weight'] = round(df['market_cap'] / df['sum'], 4)
    return df


def drop_str(str):
    str1 = str[:6]
    if str1[0] == '6':
        str1 = str1 +'.SH'
    else:
        str1 = str1 + '.SZ'
    return str1

# 记录上证50名单调整日期
'''
dateList = ['2014-06-03','2014-12-01','2015-05-14','2015-11-30','2016-05-30',
            '2016-12-28','2017-05-31','2018-05-28','2018-12-03']
'''


def count_net_north(str):
    name = str
    dateList1 = ['2014-03-31', '2014-06-03', '2014-06-30', '2014-09-30', '2014-12-01', '2014-12-31', '2015-03-31',
            '2015-05-14', '2015-06-30', '2015-09-30', '2015-11-30', '2015-12-31', '2016-03-31', '2016-05-30', '2016-06-30',
            '2016-09-30', '2016-12-28', '2016-12-31', '2017-03-31', '2017-05-31', '2017-06-30', '2017-09-30', '2017-12-31',
            '2018-03-31', '2018-05-28', '2018-06-30', '2018-09-30', '2018-12-03', '2018-12-31', '2019-03-31', '2019-06-30']
    count = get_query_count()
    print(count)

    stockData = getcode('2013-09-30')
    #stockData['jingzhi'] = 1
    dateTrade = pro.trade_cal(exchange='', start_date='20140101', end_date='20181231')
    #dateTrade2 = dateTrade[dateTrade.is_open > 0]['cal_date']
    dateTrade1 = dateTrade['cal_date']
    df = pd.DataFrame()
    x=1
    i=1
    for date in dateTrade1:
        print(date)
        df1 = pd.DataFrame()
        df2 = pd.DataFrame(columns=['date','values'])
        dateList2 = charge_time(dateList1)
        if date in dateList2:
            stockData = getcode(charge_time1(date))
            #stockData['jingzhi'] = x
        for index,row in stockData.iterrows():
            print(i)
            if i>500:
                print('一分钟请求超过500，休息一分钟继续')
                time.sleep(60)
                i = 1
            else:

                price = get_data(drop_str(row['code']), date, date)
                price['weight'] = row['weight']
               # price['jingzhi'] = row['jingzhi']
                df1 = df1.append(price)
                i = i+1
        print(df1)
        if df1.empty == False:
            df1['values'] = df1['pct_chg']*df1['weight']

            x = df1['values'].sum() / 100

            dict = {'date': date,
                    'values': x
            }

            df2 = df2.append(dict,ignore_index=True)
            df = df.append(df2)

            #stockData['jingzhi'] = x
        print(df)
    h5 = pd.HDFStore('stock_net_north'+name+'.h5', 'w', complevel=4, complib='blosc')
    h5['data'] = df
    h5.close()


count_net_north('market_cap_b')
